# pip install pycuda
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule

import numpy as np

# 定义 CUDA 内核函数
mod = SourceModule("""
__global__ void multiply_them(float *dest, float *a, float *b)
{
  const int i = threadIdx.x;
  dest[i] = a[i] * b[i];
}
""")

# 获取 CUDA 内核函数
multiply_them = mod.get_function("multiply_them")

# 生成输入数据
a = np.random.randn(400).astype(np.float32)
b = np.random.randn(400).astype(np.float32)
dest = np.zeros_like(a)

# 调用 CUDA 内核函数
multiply_them(
    cuda.Out(dest), cuda.In(a), cuda.In(b),
    block=(400, 1, 1), grid=(1, 1))

# 打印结果
print("结果:", dest)